<html>

<head>
  <meta charset="UTF-8">
  <meta name="viewport" content="width=device-width, initial-scale=1, maximum-scale=1" />
  <meta name="color-scheme" content="light dark">
  <title>llama.cpp - chat</title>

  <style>
    body {
      font-family: system-ui;
      font-size: 90%;
    }

    #container {
      margin: 0em auto;
      display: flex;
      flex-direction: column;
      justify-content: space-between;
      height: 100%;
    }

    main {
      margin: 3px;
      display: flex;
      flex-direction: column;
      justify-content: space-between;
      gap: 1em;

      flex-grow: 1;
      overflow-y: auto;

      border: 1px solid #ccc;
      border-radius: 5px;
      padding: 0.5em;
    }

    body {
      max-width: 600px;
      min-width: 300px;
      line-height: 1.2;
      margin: 0 auto;
      padding: 0 0.5em;
    }

    p {
      overflow-wrap: break-word;
      word-wrap: break-word;
      hyphens: auto;
      margin-top: 0.5em;
      margin-bottom: 0.5em;
    }

    #write form {
      margin: 1em 0 0 0;
      display: flex;
      flex-direction: column;
      gap: 0.5em;
      align-items: stretch;
    }

    .right {
      display: flex;
      flex-direction: row;
      gap: 0.5em;
      justify-content: flex-end;
    }

    fieldset {
      border: none;
      padding: 0;
      margin: 0;
    }

    fieldset.two {
      display: grid;
      grid-template: "a a";
      gap: 1em;
    }

    fieldset.three {
      display: grid;
      grid-template: "a a a";
      gap: 1em;
    }

    details {
      border: 1px solid #aaa;
      border-radius: 4px;
      padding: 0.5em 0.5em 0;
      margin-top: 0.5em;
    }

    summary {
      font-weight: bold;
      margin: -0.5em -0.5em 0;
      padding: 0.5em;
      cursor: pointer;
    }

    details[open] {
      padding: 0.5em;
    }


    textarea {
      padding: 5px;
      flex-grow: 1;
      width: 100%;
    }

    pre code {
      display: block;
      background-color: #222;
      color: #ddd;
    }
    code {
      font-family: monospace;
      padding: 0.1em 0.3em;
      border-radius: 3px;
    }

    fieldset label {
      margin: 0.5em 0;
      display: block;
    }

    header, footer {
      text-align: center;
    }

    footer {
      font-size: 80%;
      color: #888;
    }
  </style>

  <script type="module">
    import {
      html, h, signal, effect, computed, render, useSignal, useEffect, useRef
    } from '/index.js';

    import { llama } from '/completion.js';
    import { SchemaConverter } from '/json-schema-to-grammar.mjs';

    const session = signal({
      prompt: "This is a conversation between user and llama, a friendly chatbot. respond in simple markdown.",
      template: "{{prompt}}\n\n{{history}}\n{{char}}:",
      historyTemplate: "{{name}}: {{message}}",
      transcript: [],
      type: "chat",
      char: "llama",
      user: "User",
    })

    const params = signal({
      n_predict: 400,
      temperature: 0.7,
      repeat_last_n: 256, // 0 = disable penalty, -1 = context size
      repeat_penalty: 1.18, // 1.0 = disabled
      top_k: 40, // <= 0 to use vocab size
      top_p: 0.5, // 1.0 = disabled
      tfs_z: 1.0, // 1.0 = disabled
      typical_p: 1.0, // 1.0 = disabled
      presence_penalty: 0.0, // 0.0 = disabled
      frequency_penalty: 0.0, // 0.0 = disabled
      mirostat: 0, // 0/1/2
      mirostat_tau: 5, // target entropy
      mirostat_eta: 0.1, // learning rate
      grammar: '',
    })

    /* START: Support for storing prompt templates and parameters in borwser LocalStorage */

    const local_storage_storageKey = "llamacpp_server_local_storage";

    function local_storage_setDataFromObject(tag, content) {
      localStorage.setItem(local_storage_storageKey + '/' + tag, JSON.stringify(content));
    }

    function local_storage_setDataFromRawText(tag, content) {
      localStorage.setItem(local_storage_storageKey + '/' + tag, content);
    }

    function local_storage_getDataAsObject(tag) {
      const item = localStorage.getItem(local_storage_storageKey + '/' + tag);
      if (!item) {
        return null;
      } else {
        return JSON.parse(item);
      }
    }

    function local_storage_getDataAsRawText(tag) {
      const item = localStorage.getItem(local_storage_storageKey + '/' + tag);
      if (!item) {
        return null;
      } else {
        return item;
      }
    }

    // create a container for user templates and settings

    const savedUserTemplates = signal({})
    const selectedUserTemplate = signal({ name: '', template: { session: {}, params: {} } })

    // let's import locally saved templates and settings if there are any
    // user templates and settings are stored in one object
    // in form of { "templatename": "templatedata" } and { "settingstemplatename":"settingsdata" }

    console.log('Importing saved templates')

    let importedTemplates = local_storage_getDataAsObject('user_templates')

    if (importedTemplates) {
      // saved templates were successfuly imported.

      console.log('Processing saved templates and updating default template')

      //console.log(importedTemplates);
      savedUserTemplates.value = importedTemplates;

      //override default template
      savedUserTemplates.value.default = { session: session.value, params: params.value }
      local_storage_setDataFromObject('user_templates', savedUserTemplates.value)
    } else {
      // no saved templates detected.

      console.log('Initializing LocalStorage and saving default template')

      savedUserTemplates.value = { "default": { session: session.value, params: params.value } }
      local_storage_setDataFromObject('user_templates', savedUserTemplates.value)
    }

    function userTemplateResetToDefault() {
      console.log('Reseting themplate to default')
      selectedUserTemplate.value.name = 'default';
      selectedUserTemplate.value.data = savedUserTemplates.value['default'];
    }

    function userTemplateApply(t) {
      session.value = t.data.session;
      params.value = t.data.params;
    }

    function userTemplateResetToDefaultAndApply() {
      userTemplateResetToDefault()
      userTemplateApply(selectedUserTemplate.value)
    }

    function userTemplateLoadAndApplyAutosaved() {
      // get autosaved last used template
      let lastUsedTemplate = local_storage_getDataAsObject('user_templates_last')

      if (lastUsedTemplate) {

        console.log('Autosaved template found, restoring')

        selectedUserTemplate.value = lastUsedTemplate
      }
      else {

        console.log('No autosaved template found, using default template')
        // no autosaved last used template was found, so load from default.

        userTemplateResetToDefault()
      }

      console.log('Applying template')
      // and update internal data from templates

      userTemplateApply(selectedUserTemplate.value)
    }

    //console.log(savedUserTemplates.value)
    //console.log(selectedUserTemplate.value)

    function userTemplateAutosave() {
      console.log('Template Autosave...')
      if (selectedUserTemplate.value.name == 'default') {
        // we don't want to save over default template, so let's create a new one
        let newTemplateName = 'UserTemplate-' + Date.now().toString()
        let newTemplate = { 'name': newTemplateName, 'data': { 'session': session.value, 'params': params.value } }

        console.log('Saving as ' + newTemplateName)

        // save in the autosave slot
        local_storage_setDataFromObject('user_templates_last', newTemplate)

        // and load it back and apply
        userTemplateLoadAndApplyAutosaved()
      } else {
        local_storage_setDataFromObject('user_templates_last', { 'name': selectedUserTemplate.value.name, 'data': { 'session': session.value, 'params': params.value } })
      }
    }

    console.log('Checking for autosaved last used template')
    userTemplateLoadAndApplyAutosaved()

    /* END: Support for storing prompt templates and parameters in browsers LocalStorage */

    const llamaStats = signal(null)
    const controller = signal(null)

    const generating = computed(() => controller.value == null )
    const chatStarted = computed(() => session.value.transcript.length > 0)

    const transcriptUpdate = (transcript) => {
      session.value = {
        ...session.value,
        transcript
      }
    }

    // simple template replace
    const template = (str, extraSettings) => {
      let settings = session.value;
      if (extraSettings) {
        settings = { ...settings, ...extraSettings };
      }
      return String(str).replaceAll(/\{\{(.*?)\}\}/g, (_, key) => template(settings[key]));
    }

    // send message to server
    const chat = async (msg) => {
      if (controller.value) {
        console.log('already running...');
        return;
      }
      controller.value = new AbortController();

      transcriptUpdate([...session.value.transcript, ["{{user}}", msg]])

      const prompt = template(session.value.template, {
        message: msg,
        history: session.value.transcript.flatMap(([name, message]) => template(session.value.historyTemplate, {name, message})).join("\n"),
      });

      let currentMessage = '';
      const history = session.value.transcript

      const llamaParams = {
        ...params.value,
        stop: ["</s>", template("{{char}}:"), template("{{user}}:")],
      }

      for await (const chunk of llama(prompt, llamaParams, { controller: controller.value })) {
        const data = chunk.data;
        currentMessage += data.content;

        // remove leading whitespace
        currentMessage = currentMessage.replace(/^\s+/, "")

        transcriptUpdate([...history, ["{{char}}", currentMessage]])

        if (data.stop) {
          console.log("Completion finished: '", currentMessage, "', summary: ", data);
        }

        if (data.timings) {
          llamaStats.value = data.timings;
        }
      }

      controller.value = null;
    }

    function MessageInput() {
      const message = useSignal("")

      const stop = (e) => {
        e.preventDefault();
        if (controller.value) {
          controller.value.abort();
          controller.value = null;
        }
      }

      const reset = (e) => {
        stop(e);
        transcriptUpdate([]);
      }

      const submit = (e) => {
        stop(e);
        chat(message.value);
        message.value = "";
      }

      const enterSubmits = (event) => {
        if (event.which === 13 && !event.shiftKey) {
          submit(event);
        }
      }

      return html`
        <form onsubmit=${submit}>
          <div>
            <textarea type="text" rows=2 onkeypress=${enterSubmits} value="${message}" oninput=${(e) => message.value = e.target.value} placeholder="Say something..."/>
          </div>
          <div class="right">
            <button type="submit" disabled=${!generating.value} >Send</button>
            <button onclick=${stop} disabled=${generating}>Stop</button>
            <button onclick=${reset}>Reset</button>
          </div>
        </form>
      `
    }

    const ChatLog = (props) => {
      const messages = session.value.transcript;
      const container = useRef(null)

      useEffect(() => {
        // scroll to bottom (if needed)
        const parent = container.current.parentElement;
        if (parent && parent.scrollHeight <= parent.scrollTop + parent.offsetHeight + 300) {
          parent.scrollTo(0, parent.scrollHeight)
        }
      }, [messages])

      const chatLine = ([user, msg]) => {
        return html`<p key=${msg}><strong>${template(user)}:</strong> <${Markdownish} text=${template(msg)} /></p>`
      };

      return html`
        <section id="chat" ref=${container}>
          ${messages.flatMap(chatLine)}
        </section>`;
    };

    const ConfigForm = (props) => {
      const updateSession = (el) => session.value = { ...session.value, [el.target.name]: el.target.value }
      const updateParams = (el) => params.value = { ...params.value, [el.target.name]: el.target.value }
      const updateParamsFloat = (el) => params.value = { ...params.value, [el.target.name]: parseFloat(el.target.value) }
      const updateParamsInt = (el) => params.value = { ...params.value, [el.target.name]: Math.floor(parseFloat(el.target.value)) }

      const grammarJsonSchemaPropOrder = signal('')
      const updateGrammarJsonSchemaPropOrder = (el) => grammarJsonSchemaPropOrder.value = el.target.value
      const convertJSONSchemaGrammar = () => {
        try {
          const schema = JSON.parse(params.value.grammar)
          const converter = new SchemaConverter(
            grammarJsonSchemaPropOrder.value
              .split(',')
              .reduce((acc, cur, i) => ({...acc, [cur.trim()]: i}), {})
          )
          converter.visit(schema, '')
          params.value = {
            ...params.value,
            grammar: converter.formatGrammar(),
          }
        } catch (e) {
          alert(`Convert failed: ${e.message}`)
        }
      }

      const FloatField = ({label, max, min, name, step, value}) => {
        return html`
          <div>
            <label for="${name}">${label}</label>
            <input type="range" id="${name}" min="${min}" max="${max}" step="${step}" name="${name}" value="${value}" oninput=${updateParamsFloat} />
            <span>${value}</span>
          </div>
        `
      };

      const IntField = ({label, max, min, name, value}) => {
        return html`
          <div>
            <label for="${name}">${label}</label>
            <input type="range" id="${name}" min="${min}" max="${max}" name="${name}" value="${value}" oninput=${updateParamsInt} />
            <span>${value}</span>
          </div>
        `
      };

      const userTemplateReset = (e) => {
        e.preventDefault();
        userTemplateResetToDefaultAndApply()
      }

      const UserTemplateResetButton = () => {
        if (selectedUserTemplate.value.name == 'default') {
          return html`
            <button disabled>Using default template</button>
          `
        }

        return html`
          <button onclick=${userTemplateReset}>Reset all to default</button>
        `
      };

      useEffect(() => {
        // autosave template on every change
        userTemplateAutosave()
      }, [session.value, params.value])

      return html`
        <form>
          <fieldset>
            <${UserTemplateResetButton}/>
          </fieldset>

          <fieldset>
            <div>
              <label for="prompt">Prompt</label>
              <textarea type="text" name="prompt" value="${session.value.prompt}" rows=4 oninput=${updateSession}/>
            </div>
          </fieldset>

          <fieldset class="two">
            <div>
              <label for="user">User name</label>
              <input type="text" name="user" value="${session.value.user}" oninput=${updateSession} />
            </div>

            <div>
              <label for="bot">Bot name</label>
              <input type="text" name="char" value="${session.value.char}" oninput=${updateSession} />
            </div>
          </fieldset>

          <fieldset>
            <div>
              <label for="template">Prompt template</label>
              <textarea id="template" name="template" value="${session.value.template}" rows=4 oninput=${updateSession}/>
            </div>

            <div>
              <label for="template">Chat history template</label>
              <textarea id="template" name="historyTemplate" value="${session.value.historyTemplate}" rows=1 oninput=${updateSession}/>
            </div>

            <div>
              <label for="template">Grammar</label>
              <textarea id="grammar" name="grammar" placeholder="Use gbnf or JSON Schema+convert" value="${params.value.grammar}" rows=4 oninput=${updateParams}/>
              <input type="text" name="prop-order" placeholder="order: prop1,prop2,prop3" oninput=${updateGrammarJsonSchemaPropOrder} />
              <button type="button" onclick=${convertJSONSchemaGrammar}>Convert JSON Schema</button>
            </div>
          </fieldset>

          <fieldset class="two">
            ${IntField({label: "Predictions", max: 2048, min: -1, name: "n_predict", value: params.value.n_predict})}
            ${FloatField({label: "Temperature", max: 1.5, min: 0.0, name: "temperature", step: 0.01, value: params.value.temperature})}
            ${FloatField({label: "Penalize repeat sequence", max: 2.0, min: 0.0, name: "repeat_penalty", step: 0.01, value: params.value.repeat_penalty})}
            ${IntField({label: "Consider N tokens for penalize", max: 2048, min: 0, name: "repeat_last_n", value: params.value.repeat_last_n})}
            ${IntField({label: "Top-K sampling", max: 100, min: -1, name: "top_k", value: params.value.top_k})}
            ${FloatField({label: "Top-P sampling", max: 1.0, min: 0.0, name: "top_p", step: 0.01, value: params.value.top_p})}
          </fieldset>
          <details>
            <summary>More options</summary>
            <fieldset class="two">
              ${FloatField({label: "TFS-Z", max: 1.0, min: 0.0, name: "tfs_z", step: 0.01, value: params.value.tfs_z})}
              ${FloatField({label: "Typical P", max: 1.0, min: 0.0, name: "typical_p", step: 0.01, value: params.value.typical_p})}
              ${FloatField({label: "Presence penalty", max: 1.0, min: 0.0, name: "presence_penalty", step: 0.01, value: params.value.presence_penalty})}
              ${FloatField({label: "Frequency penalty", max: 1.0, min: 0.0, name: "frequency_penalty", step: 0.01, value: params.value.frequency_penalty})}
            </fieldset>
            <hr />
            <fieldset class="three">
              <div>
                <label><input type="radio" name="mirostat" value="0" checked=${params.value.mirostat == 0} oninput=${updateParamsInt} /> no Mirostat</label>
                <label><input type="radio" name="mirostat" value="1" checked=${params.value.mirostat == 1} oninput=${updateParamsInt} /> Mirostat v1</label>
                <label><input type="radio" name="mirostat" value="2" checked=${params.value.mirostat == 2} oninput=${updateParamsInt} /> Mirostat v2</label>
              </div>
              ${FloatField({label: "Mirostat tau", max: 10.0, min: 0.0, name: "mirostat_tau", step: 0.01, value: params.value.mirostat_tau})}
              ${FloatField({label: "Mirostat eta", max: 1.0, min: 0.0, name: "mirostat_eta", step: 0.01, value: params.value.mirostat_eta})}
            </fieldset>
          </details>
        </form>
      `
    }
    // poor mans markdown replacement
    const Markdownish = (params) => {
      const md = params.text
        .replace(/&/g, '&amp;')
        .replace(/</g, '&lt;')
        .replace(/>/g, '&gt;')
        .replace(/^#{1,6} (.*)$/gim, '<h3>$1</h3>')
        .replace(/\*\*(.*?)\*\*/g, '<strong>$1</strong>')
        .replace(/__(.*?)__/g, '<strong>$1</strong>')
        .replace(/\*(.*?)\*/g, '<em>$1</em>')
        .replace(/_(.*?)_/g, '<em>$1</em>')
        .replace(/```.*?\n([\s\S]*?)```/g, '<pre><code>$1</code></pre>')
        .replace(/`(.*?)`/g, '<code>$1</code>')
        .replace(/\n/gim, '<br />');
      return html`<span dangerouslySetInnerHTML=${{ __html: md }} />`;
    };

    const ModelGenerationInfo = (params) => {
      if (!llamaStats.value) {
        return html`<span/>`
      }
      return html`
        <span>
          ${llamaStats.value.predicted_per_token_ms.toFixed()}ms per token, ${llamaStats.value.predicted_per_second.toFixed(2)} tokens per second
        </span>
      `
    }

    function App(props) {

      return html`
        <div id="container">
          <header>
            <h1>llama.cpp</h1>
          </header>

          <main id="content">
            <${chatStarted.value ? ChatLog : ConfigForm} />
          </main>

          <section id="write">
            <${MessageInput} />
          </section>

          <footer>
            <p><${ModelGenerationInfo} /></p>
            <p>Powered by <a href="https://github.com/ggerganov/llama.cpp">llama.cpp</a> and <a href="https://ggml.ai">ggml.ai</a>.</p>
          </footer>
        </div>
      `;
    }

    render(h(App), document.body);
  </script>
</head>

<body>
</body>

</html>
